import streamlit as st
import time
import datetime
from dotenv import load_dotenv
from PIL import Image
from news_generator import NewsGenerator
from history_manager import HistoryManager
from document_generator import create_news_word_document
from database_manager import DatabaseManager

# 加载环境变量
load_dotenv()

# 设置页面配置
st.set_page_config(
    page_title="校园新闻自动生成系统",
    page_icon="🎓",
    layout="wide"
)

# 初始化session state
if 'user_id' not in st.session_state:
    st.session_state.user_id = None
if 'username' not in st.session_state:
    st.session_state.username = None
if 'is_logged_in' not in st.session_state:
    st.session_state.is_logged_in = False
# 添加页面状态管理
if 'page_status' not in st.session_state:
    st.session_state.page_status = 'default'
# 初始化上传图片的session state
if 'uploaded_images' not in st.session_state:
    st.session_state.uploaded_images = []
# 添加一个文件上传器重置标志
if 'reset_uploader' not in st.session_state:
    st.session_state.reset_uploader = False

# 初始化数据库管理器
db_manager = None
try:
    db_manager = DatabaseManager()
except ConnectionError as e:
    st.warning(f"数据库连接失败: {str(e)}")

# 初始化新闻生成器
try:
    news_generator = NewsGenerator()
except ValueError as e:
    # 如果API密钥未设置，创建一个模拟的新闻生成器以便演示
    class MockNewsGenerator:
        def generate_news(self, **kwargs):
            return f"模拟生成的{kwargs.get('news_type', '活动报道')}新闻\n\n根据关键词'{kwargs.get('title_keywords', '示例')}'生成的演示内容。\n\n注意：此为演示模式，请在.env文件中配置有效的DASHSCOPE_API_KEY以使用真实的生成功能。"
        
        def generate_image(self, news_content=None, title=None):
            # 模拟图片生成，返回None或一个简单的占位图
            return None
    news_generator = MockNewsGenerator()
    st.warning("未检测到有效的通义大模型API密钥，系统将以演示模式运行")

# 用户认证区域(登录和注册)
def show_login_form():
    with st.container():
        # 标题居中
        col_empty1, col_title, col_empty2 = st.columns([1, 1, 1])
        with col_title:
            st.subheader("用户登录")
        # 表单居中并限制宽度
        col1, col2, col3 = st.columns([1, 1, 1])
        with col2:
            with st.form("login_form", clear_on_submit=False):
                st.markdown("**请输入您的账号信息**")
                username = st.text_input("用户名", key="login_username", placeholder="请输入用户名")
                password = st.text_input("密码", type="password", key="login_password", placeholder="请输入密码")
                st.write(" ")
                
                # 登录表单提交
                submit = st.form_submit_button("登录", type="primary")
                
                if submit:
                    if not username or not password:
                        st.error("请输入用户名和密码")
                    else:
                        success, result = db_manager.login_user(username, password)
                        if success:
                            # 更新session状态
                            st.session_state.user_id = result['id']
                            st.session_state.username = result['username']
                            st.session_state.is_logged_in = True
                            # 设置页面状态为登录成功
                            st.session_state.page_status = 'login_success'
                            # 强制页面重新运行
                            st.rerun()
                        else:
                            # 显示具体的错误信息
                            if result == "用户名不存在":
                                st.error("该用户不存在，请检查输入")
                            elif result == "密码错误":
                                st.error("密码错误，请重新输入")
                            else:
                                st.error(result)

def show_register_form():
    with st.container():
        # 使用列布局让标题居中
        col_empty1, col_title, col_empty2 = st.columns([1, 1, 1])
        with col_title:
            st.subheader("用户注册")
        col1, col2, col3 = st.columns([1, 1, 1])
        with col2:
            with st.form("register_form"):
                st.markdown("**创建您的账号**")
                username = st.text_input("用户名", key="register_username", placeholder="请输入用户名")
                password = st.text_input("密码", type="password", key="register_password", placeholder="请设置密码（6位及以上）")
                confirm_password = st.text_input("确认密码", type="password", key="confirm_password", placeholder="请再次输入密码")
                
                # 注册表单提交
                submit_register = st.form_submit_button("注册", type="primary")
                
                if submit_register:
                    if not username or not password:
                        st.error("用户名和密码不能为空")
                    elif password != confirm_password:
                        st.error("两次输入的密码不一致")
                    elif len(password) < 6:
                        st.error("密码长度不能少于6位")
                    else:
                        success, message = db_manager.register_user(username, password)
                        if success:
                            st.success(message)
                            st.info("请使用新账号登录")
                        else:
                            st.error(message)

def show_user_info():
    col1, col2 = st.columns([1, 1])
    with col1:
        st.write(f"👤 当前用户: **{st.session_state.username}**")
    with col2:
        if st.button("退出登录", type="secondary"):
            # 清除session状态
            st.session_state.user_id = None
            st.session_state.username = None
            st.session_state.is_logged_in = False
            # 设置页面状态为退出成功
            st.session_state.page_status = 'logout_success'
            # 强制页面重新运行
            st.rerun()

# 页面状态处理
if st.session_state.page_status == 'login_success':
    # 显示登录成功信息并延迟跳转
    st.success("登录成功！正在跳转到主页面...")
    time.sleep(1)  # 给用户足够时间看到成功消息
    st.session_state.page_status = 'default'  # 重置页面状态
    st.rerun()  # 重新运行以加载主界面

if st.session_state.page_status == 'logout_success':
    # 显示退出成功信息并延迟跳转
    st.success("已退出登录！")
    time.sleep(1)  # 给用户足够时间看到成功消息
    st.session_state.page_status = 'default'  # 重置页面状态
    st.rerun()  # 重新运行以加载登录界面

# 显示认证界面或用户信息管理
if not st.session_state.is_logged_in:
    if db_manager is None:
        st.error("数据库连接失败，无法登录。请联系管理员。")
        st.stop()
    else:
        # 显示登录/注册界面
        st.title("🎓 校园新闻自动生成系统")
        tab1_auth, tab2_auth = st.tabs(["登录", "注册"])
        with tab1_auth:
            show_login_form()
        with tab2_auth:
            show_register_form()
        st.stop()
else:
    # 用户已登录，即使 db_manager 为 None 也继续
    st.title("🎓 校园新闻自动生成系统")
    show_user_info()

# 初始化历史记录管理器
history_manager = HistoryManager(user_id=st.session_state.user_id)

# 创建标签页
tab1, tab2 = st.tabs(["生成新闻", "历史记录"])

with tab1:
    # 侧边栏配置
    with st.sidebar:
        st.header("配置选项")
        
        # 生成新闻类型选择
        news_type = st.selectbox(
            "选择新闻类型",
            ["活动报道", "公告通知", "人物专访", "学术动态", "校园生活", "其他"]
        )
        
        # 生成长度控制
        max_length = st.slider(
            "新闻长度（字数）",
            min_value=300,
            max_value=2000,
            value=800,
            step=100
        )
        
        # 语言风格选择
        language_style = st.selectbox(
            "语言风格",
            ["正式专业", "通俗易懂", "生动活泼", "简洁明快", "详细全面"]
        )
        
        # 生成温度控制（影响创造性）
        temperature = st.slider(
            "创造性程度",
            min_value=0.1,
            max_value=1.0,
            value=0.7,
            step=0.1
        )
        
        # 是否生成配图
        is_generate_image = st.checkbox(
            "生成新闻AI配图",
            value=False,
            help="自动生成相关的AI配图并融入新闻内容"
        )
        
        # 模型选择
        model = st.selectbox(
            "选择模型",
            ["qwen-turbo", "qwen-plus", "qwen-max"]
        )

    # 主要内容输入区域
    st.subheader("新闻内容参数设置")

    # 创建两列布局
    col1, col2 = st.columns(2)

    with col1:
        # 新闻标题关键词
        title_keywords = st.text_input(
            "新闻标题关键词",
            placeholder="例如：毕业典礼、科技创新大赛、校长讲话..."
        )
        
        # 主要人物
        key_persons = st.text_input(
            "主要人物(选填)",
            placeholder="例如：张三校长、李四教授、王五同学..."
        )
        
        # 时间地点
        time_location = st.text_input(
            "时间地点(选填)",
            placeholder="例如：2023年6月18日、学校体育馆..."
        )

    with col2:
        # 核心内容要点
        content_points = st.text_area(
            "核心内容要点",
            placeholder="请列出新闻的关键信息点，每行一个要点...",
            height=150
        )
        
        # 附加要求
        additional_requirements = st.text_area(
            "附加要求(选填)",
            placeholder="如有其他特殊要求，请在此说明...",
            height=80
        )
    
    # 图片上传区域
    st.subheader("上传图片")
    # 自定义图片上传区域样式
    st.markdown(
    """<style>
        /* 自定义文件上传区域提示文本 */
        [data-testid="stFileUploaderDropzone"] div div::before {
        content: "可拖放文件到此处";
        }

        /* 隐藏原始拖拽提示文本 */
        [data-testid="stFileUploaderDropzone"] div div span {
            display: none;
        }

        /* 隐藏原始限制提示文本 */
        [data-testid="stFileUploaderDropzone"] div div small {
            display: none;
        }       

        /* 添加文件大小限制提示 */
        [data-testid="stFileUploaderDropzone"] div div::after {
            color: rgba(49, 51, 63, 0.6);
            font-size: 0.8em;
            content: "每张图片限制小于15MB";
        }   
    </style>""", unsafe_allow_html=True)
    
    # 使用动态key来重置文件上传器
    uploader_key = "file_uploader" + ("_reset" if st.session_state.reset_uploader else "")
    
    uploaded_files = st.file_uploader(
        "上传相关图片（可选），支持多张，单张不超过15MB", 
        type=["jpg", "jpeg", "png"], 
        accept_multiple_files=True,
        help="请上传JPG、JPEG或PNG格式的图片文件，单张图片不超过15MB",
        key=uploader_key,
        # 当reset_uploader为True时，强制清除文件选择
        on_change=lambda: setattr(st.session_state, "reset_uploader", False) if st.session_state.reset_uploader else None
    )
    
    # 处理图片大小验证并更新session state
    max_size_mb = 15
    max_size_bytes = max_size_mb * 1024 * 1024
    
    if uploaded_files:
        # 重置已上传图片列表
        st.session_state.uploaded_images = []
        for file in uploaded_files:
            # 检查文件大小
            file.seek(0, 2)  # 移动到文件末尾
            file_size = file.tell()
            file.seek(0)  # 重置文件指针
            if file_size > max_size_bytes:
                st.error(f"图片 '{file.name}' 文件大小 {file_size/1024/1024:.2f}MB 超过 {max_size_mb}MB 上限，请移除")
            else:
                st.session_state.uploaded_images.append(file)
    
    # 显示上传的图片预览和清空按钮
    if st.session_state.uploaded_images:
        # 显示图片数量和清空按钮
        col_info, col_clear = st.columns([3, 1])
        with col_info:
            st.write(f"已成功上传 {len(st.session_state.uploaded_images)} 张图片")
        with col_clear:
            if st.button("清空图片", type="secondary", use_container_width=True):
                # 清空session state中的图片列表
                st.session_state.uploaded_images = []
                # 设置重置标志，这将改变file_uploader的key
                st.session_state.reset_uploader = True
                # 强制页面重新加载
                st.rerun()
        
        # 使用网格显示图片
        num_images = len(st.session_state.uploaded_images)
        cols = 2  # 每行显示2张图片
        rows = (num_images + cols - 1) // cols
        
        for row in range(rows):
            cols_to_use = st.columns(cols)
            for col_idx in range(cols):
                img_idx = row * cols + col_idx
                if img_idx < num_images:
                    with cols_to_use[col_idx]:
                        image = Image.open(st.session_state.uploaded_images[img_idx])
                        st.image(image, caption=f"图片 {img_idx + 1}")

    # 生成按钮
    if st.button("生成新闻", type="primary"):
        # 表单验证
        if not title_keywords or not content_points:
            st.error("请至少填写标题关键词和核心内容要点")
        else:
            with st.spinner("正在生成新闻..."):
                try:
                    # 准备图片列表和分析结果
                    uploaded_images_list = []
                    image_analyses = []
                    
                    # 检查用户是否上传了图片，如果有则分析
                    if st.session_state.uploaded_images:
                        for idx, uploaded_img in enumerate(st.session_state.uploaded_images):
                            img = Image.open(uploaded_img)
                            uploaded_images_list.append(img)
                            # 分析图片信息，用于补充新闻生成
                            try:
                                img_analysis = news_generator.analyze_image(img)
                                image_analyses.append(f"图片{idx+1}分析结果: {img_analysis}")
                            except Exception as e:
                                st.warning(f"分析图片 {idx+1} 时出错: {str(e)}")
                                image_analyses.append("无法分析图片内容")
                    
                    # 生成新闻内容，同时提取标题
                    # 将图片分析结果传递给新闻生成器
                    news_content, news_title = news_generator.generate_news(
                        news_type=news_type,
                        title_keywords=title_keywords,
                        key_persons=key_persons,
                        time_location=time_location,
                        content_points=content_points,
                        language_style=language_style,
                        max_length=max_length,
                        additional_requirements=additional_requirements,
                        temperature=temperature,
                        model=model,
                        image_analyses=image_analyses
                    )
                    
                    generated_img = None
                    # 如果用户选择生成配图，调用图像生成功能
                    if is_generate_image:
                        with st.spinner("正在生成AI配图..."):
                            try:
                                generated_img = news_generator.generate_image(news_content, news_title)
                            except Exception as e:
                                st.warning(f"AI配图生成失败: {str(e)}")

                    # 准备用户上传的图片列表，包含描述
                    user_images_with_desc = []
                    # 添加用户上传的图片
                    if uploaded_images_list:
                        for idx, img in enumerate(uploaded_images_list):
                            try:
                                # 使用之前的分析结果或重新分析
                                uploaded_img_description = image_analyses[idx] if idx < len(image_analyses) else ""
                                if not uploaded_img_description or uploaded_img_description == "无法分析图片内容":
                                    uploaded_img_description = news_generator.analyze_image(img)
                                # 提取简短描述
                                user_images_with_desc.append((img, f"图片{idx+1}分析结果:{uploaded_img_description}"))
                            except Exception as e:
                                # 如果分析失败，使用默认描述
                                user_images_with_desc.append((img, f"用户上传的相关图片{idx+1}"))
                    
                    # 准备下载文件的元数据 
                    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                    txt_filename = f"campus_news_{timestamp}.txt"
                    doc_filename = f"campus_news_{timestamp}.docx"

                    # 准备下载Word文档，按新要求插入图片
                    doc_bytes = create_news_word_document(
                        news_title,
                        news_content,
                        user_images=user_images_with_desc,
                        ai_image=generated_img
                    )
                    
                    # 保存到历史记录，同时保存文档
                    history_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    
                    st.session_state.generated_news = {
                        'title': news_title,
                        'txt_content': news_content,
                        'doc_content': doc_bytes,
                        'txt_filename': txt_filename,
                        'doc_filename': doc_filename,
                        'history_time': history_time
                    }
                    
                    # 保存到历史记录
                    history_manager.add_record(
                        title=news_title,
                        title_keywords=title_keywords,
                        news_type=news_type,
                        txt_content=news_content,
                        doc_filename=doc_filename,
                        doc_content=doc_bytes,
                        created_at=history_time
                    )

                    # 输出生成的新闻内容供用户查看
                    st.subheader("📝 生成内容预览(仅文本)")
                    st.markdown(f"**标题：** {news_title}")
                    st.markdown(news_content)

                    # 添加下载功能 - 纯文本和Word文档格式选择
                    st.subheader("💾 下载")
                    if 'generated_news' in st.session_state:
                        news_data = st.session_state.generated_news
                        col1, col2 = st.columns(2)
                        with col1:
                            st.download_button(
                                label="下载纯文本文件",
                                data=news_data['txt_content'],
                                file_name=news_data['txt_filename'],
                                mime="text/plain",  # 文件MIME类型，表示纯文本
                                use_container_width=True,
                                help="下载纯文本格式的新闻稿"
                            )

                        with col2:
                            st.download_button(
                                label="下载Word文档",
                                data=news_data['doc_content'],
                                file_name=news_data['doc_filename'],
                                mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document",  # Word文档的MIME类型
                                use_container_width=True,
                                help="下载(包含图片)的Word文档格式新闻稿"
                            )
                except Exception as e:
                    st.error(f"生成过程中出错: {str(e)}")
with tab2:
    st.subheader("历史记录")
    # 显示历史记录
    records = history_manager.get_records()
    if not records:
        st.info("暂无历史记录")
    else:
        # 添加清空历史记录按钮
        col_clear, col_empty = st.columns([3, 5])
        with col_clear:
            if st.button("清空历史记录", type="secondary"):
                all_clear_success = history_manager.clear_history()
                if all_clear_success:
                    st.rerun()
                else:
                    st.error("清空历史记录失败，请重试")
                    
        # 显示记录卡片
        for i, record in enumerate(records):
            # 如果有完整标题则使用标题，否则回退到关键词
            display_title = record.get('title', record.get('title_keywords', '无标题'))
            with st.expander(f"[{record['news_type']}] {display_title} - {record['timestamp']}"):
                st.markdown(record['txt_content'])
                
                # 创建操作按钮行
                col1, col2, col3 = st.columns([1, 1, 1])
                # 创建一行显示三个按钮，让它们相邻挨着
                buttons_col = st.columns([1, 1, 1], gap='small')
                # 下载纯文本文件按钮
                with buttons_col[0]:
                    try:
                        # 尝试使用新方法获取文本内容
                        txt_record = history_manager.get_document(record.get('record_id', ''))
                        if txt_record:
                            download_txt_filename = txt_record['doc_filename'][:-5] + '.txt'
                            st.download_button(
                                label="下载文本文件",
                                data=txt_record['txt_content'],
                                file_name=download_txt_filename,
                                mime="text/plain",
                                key=f"download_txt_{txt_record.get('record_id', i)}",
                                use_container_width=True
                            )
                    except Exception as e:
                        st.warning(f"无法获取文本文件：{str(e)}")
                
                # 下载Word文档按钮
                with buttons_col[1]:
                    try:
                        doc_record = history_manager.get_document(record.get('record_id', ''))
                        if doc_record:
                            st.download_button(
                                label="下载Word文档",
                                data=doc_record['doc_content'],
                                file_name=doc_record['doc_filename'],
                                mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
                                key=f"download_doc_{doc_record.get('record_id', i)}",
                                use_container_width=True
                            )
                    except Exception as e:
                        st.warning(f"无法获取Word文档：{str(e)}")
                
                # 删除记录按钮
                with buttons_col[2]:
                    if st.button(
                        "删除记录", 
                        type="secondary", 
                        key=f"delete_{record.get('record_id', i)}",
                        help="删除此条历史记录",
                        use_container_width=True
                    ):
                        success = history_manager.delete_record(record.get('record_id', ''))
                        if success:
                            st.rerun()
                        else:
                            st.error("删除失败，请重试")

# 页脚
st.markdown("---")
st.markdown("校园新闻自动生成系统 © 2025")